
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include "../utils/BaseConvert.h"
#include "../utils/StringUtils.h"
#include "../iface/debug.h"
#include "RSAKeyPair.h"

int RSAKeyPair::InitialiseRSAKeyFromPublicKey(CRYPT_PKCINFO_RSA *rsaKey)
{
   unsigned char *buffera[2];
   cryptInitComponents(rsaKey, CRYPT_KEYTYPE_PUBLIC);
   buffera[0] = (unsigned char *)malloc(BN_num_bytes(pub->n));
   if (!(buffera[0])) { errprint("Could not allocate buffer for pub->n."); return -1; }
   buffera[1] = (unsigned char *)malloc(BN_num_bytes(pub->e));
   if (!(buffera[1])) { errprint("Could not allocate buffer for pub->e."); return -1; }
   BN_bn2bin(pub->n, buffera[0]);
   BN_bn2bin(pub->e, buffera[1]);
   cryptSetComponent(rsaKey->n, buffera[0], BN_num_bits(pub->n));
   cryptSetComponent(rsaKey->e, buffera[1], BN_num_bits(pub->e));
   for (int i = 0; i < 2; i++) free(buffera[i]);
   return 0;
}

int RSAKeyPair::InitialiseRSAKeyFromPrivateKey(CRYPT_PKCINFO_RSA *rsaKey)
{
   unsigned char *buffera[8];
   cryptInitComponents(rsaKey, CRYPT_KEYTYPE_PRIVATE);
   buffera[0] = (unsigned char *)malloc(BN_num_bytes(prv->n));
   if (!(buffera[0])) { errprint("Could not allocate buffer for prv->n."); return -1; }
   buffera[1] = (unsigned char *)malloc(BN_num_bytes(prv->e));
   if (!(buffera[1])) { errprint("Could not allocate buffer for prv->e."); return -1; }
   buffera[2] = (unsigned char *)malloc(BN_num_bytes(prv->d));
   if (!(buffera[2])) { errprint("Could not allocate buffer for prv->d."); return -1; }
   buffera[3] = (unsigned char *)malloc(BN_num_bytes(prv->p));
   if (!(buffera[3])) { errprint("Could not allocate buffer for prv->p."); return -1; }
   buffera[4] = (unsigned char *)malloc(BN_num_bytes(prv->q));
   if (!(buffera[4])) { errprint("Could not allocate buffer for prv->q."); return -1; }
   buffera[5] = (unsigned char *)malloc(BN_num_bytes(prv->iqmp));
   if (!(buffera[5])) { errprint("Could not allocate buffer for prv->iqmp."); return -1; }
   buffera[6] = (unsigned char *)malloc(BN_num_bytes(prv->dmp1));
   if (!(buffera[6])) { errprint("Could not allocate buffer for prv->dmp1."); return -1; }
   buffera[7] = (unsigned char *)malloc(BN_num_bytes(prv->dmq1));
   if (!(buffera[7])) { errprint("Could not allocate buffer for prv->dmq1."); return -1; }
   BN_bn2bin(prv->n, buffera[0]);
   BN_bn2bin(prv->e, buffera[1]);
   BN_bn2bin(prv->d, buffera[2]);
   BN_bn2bin(prv->p, buffera[3]);
   BN_bn2bin(prv->q, buffera[4]);
   BN_bn2bin(prv->iqmp, buffera[5]);
   BN_bn2bin(prv->dmp1, buffera[6]);
   BN_bn2bin(prv->dmq1, buffera[7]);
   cryptSetComponent(rsaKey->n, buffera[0], BN_num_bits(prv->n));
   cryptSetComponent(rsaKey->e, buffera[1], BN_num_bits(prv->e));
   cryptSetComponent(rsaKey->d, buffera[2], BN_num_bits(prv->d));
   cryptSetComponent(rsaKey->p, buffera[3], BN_num_bits(prv->p));
   cryptSetComponent(rsaKey->q, buffera[4], BN_num_bits(prv->q));
   cryptSetComponent(rsaKey->u, buffera[5], BN_num_bits(prv->iqmp));
   cryptSetComponent(rsaKey->e1, buffera[6], BN_num_bits(prv->dmp1));
   cryptSetComponent(rsaKey->e2, buffera[7], BN_num_bits(prv->dmq1));
   for (int i = 0; i < 8; i++) free(buffera[i]);
   return 0;
}

CRYPT_CONTEXT RSAKeyPair::GetCryptContext_PublicKey(const char *name)
{
   int status;
   CRYPT_CONTEXT cntxt;
   CRYPT_PKCINFO_RSA *rsaKey = (CRYPT_PKCINFO_RSA *)malloc(sizeof(CRYPT_PKCINFO_RSA));
   InitialiseRSAKeyFromPublicKey(rsaKey);
   int magictriescntxt = 0;
   do
   {
      status = cryptCreateContext(&cntxt, CRYPT_UNUSED, CRYPT_ALGO_RSA);
      if (status == CRYPT_ERROR_TIMEOUT) apr_sleep(rand() % CRYPT_MAX_TIMESLEEP);
      magictriescntxt++;
   } while ((status == CRYPT_ERROR_TIMEOUT) && (magictriescntxt < CRYPT_MAX_RETRIES));
   fCryptCheckError(status);
   status = cryptSetAttributeString(cntxt, CRYPT_CTXINFO_LABEL, name, strlen(name));
   fCryptCheckError(status);
   status = cryptSetAttributeString(cntxt, CRYPT_CTXINFO_KEY_COMPONENTS, rsaKey, sizeof(CRYPT_PKCINFO_RSA));
   fCryptCheckError(status);
   cryptDestroyComponents(rsaKey);
   free(rsaKey);
   return cntxt;
}

CRYPT_CONTEXT RSAKeyPair::GetCryptContext_PrivateKey(const char *name)
{
   int status;
   CRYPT_CONTEXT cntxt;
   CRYPT_PKCINFO_RSA *rsaKey = (CRYPT_PKCINFO_RSA *)malloc(sizeof(CRYPT_PKCINFO_RSA));
   InitialiseRSAKeyFromPrivateKey(rsaKey);
   int magictriescntxt = 0;
   do
   {
      status = cryptCreateContext(&cntxt, CRYPT_UNUSED, CRYPT_ALGO_RSA);
      if (status == CRYPT_ERROR_TIMEOUT) apr_sleep(rand() % CRYPT_MAX_TIMESLEEP);
      magictriescntxt++;
   } while ((status == CRYPT_ERROR_TIMEOUT) && (magictriescntxt < CRYPT_MAX_RETRIES));
   fCryptCheckError(status);
   status = cryptSetAttributeString(cntxt, CRYPT_CTXINFO_LABEL, name, strlen(name));
   fCryptCheckError(status);
   status = cryptSetAttributeString(cntxt, CRYPT_CTXINFO_KEY_COMPONENTS, rsaKey, sizeof(CRYPT_PKCINFO_RSA));
   fCryptCheckError(status);
   cryptDestroyComponents(rsaKey);
   free(rsaKey);
   return cntxt;
}

int RSAKeyPair::GenerateKeyPair(int nbits)
{
   RSA *rkey = RSA_generate_key(nbits, 0x10001, NULL, NULL);
   if (!rkey) { fSSLCheckError(); return -1; }
   RSACopyPublicKey(pub, rkey);
   RSACopyPrivateKey(prv, rkey);
   RSA_free(rkey);
   return 0;
}

RSAKeyPair::RSAKeyPair()
{
   pub = RSA_new();
   prv = RSA_new();
}

RSAKeyPair::~RSAKeyPair()
{
   RSA_free(pub);
   RSA_free(prv);
}

char *RSAKeyPair::ExportPrivateKey()
{
   char *data;
   int damt;
   char *dbuf = NULL;
   EVP_PKEY *epkey = EVP_PKEY_new();
   EVP_PKEY_set1_RSA(epkey, prv);
   BIO *bo = BIO_new(BIO_s_mem());
   if (!PEM_write_bio_PrivateKey(bo, epkey, 0, 0, 0, 0, 0))
   {
      errprint("could not export private key");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   damt = BIO_get_mem_data(bo, &data);
   dbuf = (char *)malloc(sizeof(char)*(damt+1));
   memcpy(dbuf, data, sizeof(char)*(damt+1));
   dbuf[damt] = '\0';
//   printf("%s\n", dbuf);
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return dbuf;
}

char *RSAKeyPair::ExportPasswordedPrivateKey(const char *password)
{
   char *data;
   int damt;
   char *dbuf = NULL;
   EVP_PKEY *epkey = EVP_PKEY_new();
   EVP_PKEY_set1_RSA(epkey, prv);
   BIO *bo = BIO_new(BIO_s_mem());
   if (!PEM_write_bio_PKCS8PrivateKey(bo, epkey, EVP_des_ede3_cbc(), NULL, 0, 0, (void *)password))
   {
      errprint("could not export private key (password optional)");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   damt = BIO_get_mem_data(bo, &data);
   dbuf = (char *)malloc(sizeof(char)*(damt+1));
   memcpy(dbuf, data, sizeof(char)*(damt+1));
   dbuf[damt] = '\0';
//   printf("%s\n", dbuf);
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return dbuf;
}

int RSAKeyPair::ImportPrivateKey(const char *encoded)
{
   return ImportPasswordedPrivateKey(NULL, encoded);
}

int RSAKeyPair::ImportPasswordedPrivateKey(const char *password, const char *encoded)
{
   int retval = 0;
   RSA *temp_key = NULL;
   BIO *bo = BIO_new_mem_buf((void *)encoded, -1);
   EVP_PKEY *epkey = PEM_read_bio_PrivateKey(bo, NULL, 0, (void *)password);
   if (!epkey)
   {
      retval = 1;
      errprint("could not import private key (password optional)");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   RSA_free(prv);
   prv = RSA_new();
   temp_key = EVP_PKEY_get1_RSA(epkey);
   if (temp_key)
   {
      RSACopyPrivateKey(prv, temp_key);
      RSA_free(temp_key);
   }
   else
      errprint("Failed private key import.");
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return retval;
}

char *RSAKeyPair::ExportPublicKey()
{
   char *data;
   int damt;
   char *dbuf = NULL;
   EVP_PKEY *epkey = EVP_PKEY_new();
   EVP_PKEY_set1_RSA(epkey, pub);
   BIO *bo = BIO_new(BIO_s_mem());
   if (!PEM_write_bio_PUBKEY(bo, epkey))
   {
      errprint("could not export public key");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   damt = BIO_get_mem_data(bo, &data);
   dbuf = (char *)malloc(sizeof(char)*(damt+1));
   memcpy(dbuf, data, sizeof(char)*(damt+1));
   dbuf[damt] = '\0';
//   printf("%s\n", dbuf);
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return dbuf;
}

char *RSAKeyPair::ExportOpenSSHPublicKeyBlob(int *reallength)
{
   const unsigned char keyheader[11] = { 0x00, 0x00, 0x00, 0x07, 0x73, 0x73, 0x68, 0x2D, 0x72, 0x73, 0x61 };
   unsigned char *n_buffer = NULL;
   unsigned char *e_buffer = NULL;
   int n_len = BN_num_bytes(pub->n);
   int e_len = BN_num_bytes(pub->e);
   n_buffer = (unsigned char *)malloc(n_len);
   if (!n_buffer) { errprint("Could not allocate buffer for pub->n."); return NULL; }
   e_buffer = (unsigned char *)malloc(e_len);
   if (!e_buffer) { errprint("Could not allocate buffer for pub->e."); return NULL; }
   BN_bn2bin(pub->n, n_buffer);
   BN_bn2bin(pub->e, e_buffer);
   int enclen = 4 + 7 + 4 + ((e_buffer[0] & 0x80) == 0x80) + e_len + 4 + ((n_buffer[0] & 0x80) == 0x80) + n_len;
   char *encode = (char *)malloc(sizeof(char)*enclen);
   memcpy(encode, keyheader, sizeof(char)*11);
   int real_len, data_start, toffset;
   toffset = 11;
   real_len = e_len + ((e_buffer[0] & 0x80) == 0x80);
   data_start = 4 + ((e_buffer[0] & 0x80) == 0x80);
   memset(&(encode[toffset]), 0, sizeof(char)*(real_len + 4));
   encode[toffset + 0] = (unsigned char)(real_len >> 24);
   encode[toffset + 1] = (unsigned char)(real_len >> 16);
   encode[toffset + 2] = (unsigned char)(real_len >>  8);
   encode[toffset + 3] = (unsigned char)(real_len);
   memcpy(&(encode[toffset + data_start]), e_buffer, sizeof(char)*e_len);
   toffset += (real_len + 4);
   real_len = n_len + ((n_buffer[0] & 0x80) == 0x80);
   data_start = 4 + ((n_buffer[0] & 0x80) == 0x80);
   memset(&(encode[toffset]), 0, sizeof(char)*(real_len + 4));
   encode[toffset + 0] = (unsigned char)(real_len >> 24);
   encode[toffset + 1] = (unsigned char)(real_len >> 16);
   encode[toffset + 2] = (unsigned char)(real_len >>  8);
   encode[toffset + 3] = (unsigned char)(real_len);
   memcpy(&(encode[toffset + data_start]), n_buffer, sizeof(char)*n_len);
   free(n_buffer); free(e_buffer);
   if (reallength) *reallength = enclen;
   return encode;
}

char *RSAKeyPair::ExportOpenSSHPublicKey(const char *keylabel)
{
   int enclen = 0;
   char *encode = ExportOpenSSHPublicKeyBlob(&enclen);
   char *base64string = BaseConvert::ConvertBytesToBase64(encode, enclen);
   free(encode);
   char *finalstring = StringUtils::varprintf("ssh-rsa %s %s", base64string, keylabel);
   free(base64string);
   return finalstring;
}

int RSAKeyPair::ImportPublicKey(const char *encoded)
{
   int retval = 0;
   RSA *temp_key = NULL;
   BIO *bo = BIO_new_mem_buf((void *)encoded, -1);
   EVP_PKEY *epkey = PEM_read_bio_PUBKEY(bo, NULL, 0, NULL);
   fSSLCheckError();
   if (!epkey)
   {
      retval = 1;
      errprint("could not import public key");
      goto bioshutdown;
   }
   RSA_free(pub);
   pub = RSA_new();
   temp_key = EVP_PKEY_get1_RSA(epkey);
   if (temp_key)
   {
      RSACopyPublicKey(pub, temp_key);
      RSA_free(temp_key);
   }
   else
      errprint("Failed public key import.");
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return retval;
}

const char *RSAKeyPair::GetKeyTypeName()
{
   return "RSA";
}

char *RSAKeyPair::GetMD5KeyFingerprint()
{
   char *outputstring = (char *)malloc(sizeof(char)*(EVP_MAX_MD_SIZE+1)*3);
   if (outputstring == NULL) return NULL;
   int nlen = BN_num_bytes(pub->n);
   int elen = BN_num_bytes(pub->e);
   int len = nlen + elen;
   unsigned char *blob = (unsigned char *)malloc(sizeof(unsigned char)*len);
   BN_bn2bin(pub->n, blob);
   BN_bn2bin(pub->e, &(blob[nlen]));
   outputstring[0] = '\0';
   if (blob != NULL)
   {
      int i;
      unsigned char digest[EVP_MAX_MD_SIZE];
      EVP_MD_CTX *ctx;
      ctx = EVP_MD_CTX_create();
      const EVP_MD *md5 = EVP_md5();
      EVP_DigestInit_ex(ctx, md5, NULL);
      EVP_DigestUpdate(ctx, blob, len);
      EVP_DigestFinal_ex(ctx, digest, NULL);
      EVP_MD_CTX_destroy(ctx);
      for (i = 0; i < md5->md_size; i++)
      {
         char hex[4];
         sprintf(hex, "%02x:", digest[i]);
         strcat(outputstring, hex);
      }
      fSSLCheckError();
      // chop last :
      outputstring[strlen(outputstring) - 1] = '\0';
      // erase key blob securely
      memset(blob, 0, len);
      free(blob);
   } else
   {
      free(outputstring);
      outputstring = NULL;
   }
   return outputstring;
}

char *RSAKeyPair::GetFingerprintAsVariableName()
{
   char *outputstring = (char *)malloc(sizeof(char)*((EVP_MAX_MD_SIZE+1)*2+1));
   if (outputstring == NULL) return NULL;
   int nlen = BN_num_bytes(pub->n);
   int elen = BN_num_bytes(pub->e);
   int len = nlen + elen;
   unsigned char *blob = (unsigned char *)malloc(sizeof(unsigned char)*len);
   BN_bn2bin(pub->n, blob);
   BN_bn2bin(pub->e, &(blob[nlen]));
   outputstring[0] = '_';
   outputstring[1] = '\0';
   if (blob != NULL)
   {
      int i;
      unsigned char digest[EVP_MAX_MD_SIZE];
      EVP_MD_CTX *ctx;
      ctx = EVP_MD_CTX_create();
      const EVP_MD *md5 = EVP_md5();
      EVP_DigestInit_ex(ctx, md5, NULL);
      EVP_DigestUpdate(ctx, blob, len);
      EVP_DigestFinal_ex(ctx, digest, NULL);
      EVP_MD_CTX_destroy(ctx);
      for (i = 0; i < md5->md_size; i++)
      {
         char hex[3];
         sprintf(hex, "%02x", digest[i]);
         strcat(outputstring, hex);
      }
      fSSLCheckError();
      // erase key blob securely
      memset(blob, 0, len);
      free(blob);
   } else
   {
      free(outputstring);
      outputstring = NULL;
   }
   return outputstring;
}

int RSAKeyPair::RSACopyPublicKey(RSA *dst, RSA *src)
{
   if (dst == NULL) { errprint("NULL public key destination"); return -1; }
   if (src == NULL) { errprint("NULL public key source"); return 1; }
   // Copy public key
   dst->n = BN_new();
   if (BN_copy(dst->n, src->n) != (dst->n)) { errprint("Could not copy pub->n"); return -1; }
   dst->e = BN_new();
   if (BN_copy(dst->e, src->e) != (dst->e)) { errprint("Could not copy pub->e"); return -1; }
   return 0;
}

int RSAKeyPair::RSACopyPrivateKey(RSA *dst, RSA *src)
{
   if (dst == NULL) { errprint("NULL private key destination"); return -1; }
   if (src == NULL) { errprint("NULL private key source"); return 1; }
   // Copy public key
   dst->n = BN_new();
   if (BN_copy(dst->n, src->n) != (dst->n)) { errprint("Could not copy prv->n"); return -1; }
   dst->e = BN_new();
   if (BN_copy(dst->e, src->e) != (dst->e)) { errprint("Could not copy prv->e"); return -1; }

   // Copy private key
   dst->d = BN_new();
   if (BN_copy(dst->d, src->d) != (dst->d)) { errprint("Could not copy prv->d"); return -1; }
   dst->p = BN_new();
   if (BN_copy(dst->p, src->p) != (dst->p)) { errprint("Could not copy prv->p"); return -1; }
   dst->q = BN_new();
   if (BN_copy(dst->q, src->q) != (dst->q)) { errprint("Could not copy prv->q"); return -1; }

   dst->dmp1 = BN_new();
   if (BN_copy(dst->dmp1, src->dmp1) != (dst->dmp1)) { errprint("Could not copy prv->dmp1"); return -1; }
   dst->dmq1 = BN_new();
   if (BN_copy(dst->dmq1, src->dmq1) != (dst->dmq1)) { errprint("Could not copy prv->dmq1"); return -1; }
   dst->iqmp = BN_new();
   if (BN_copy(dst->iqmp, src->iqmp) != (dst->iqmp)) { errprint("Could not copy prv->iqmp"); return -1; }
   return 0;
}

int RSAKeyPair::RSAKeyCompareThisPublicToPublic(RSAKeyPair *icmp)
{
   if (icmp == NULL)
   {
      errprint("comparing against NULL key");
      return 0;
   }
   EVP_PKEY *evp1 = EVP_PKEY_new();
   EVP_PKEY *evp2 = EVP_PKEY_new();
   EVP_PKEY_set1_RSA(evp1, pub);
   EVP_PKEY_set1_RSA(evp2, icmp->pub);
   int retval = EVP_PKEY_cmp(evp1, evp2);
   EVP_PKEY_free(evp1);
   EVP_PKEY_free(evp2);
   fSSLCheckError();
   switch (retval)
   {
      case -2:
         errprint("key compare operation not supported (EVP_PKEY_cmp())");
         break;
      case -1:
         errprint("tried to compare different key types (EVP_PKEY_cmp())");
         return 0;
      case 0:
         return 0;
      case 1:
         return 1;
   }
   return 0;
}

void RSAKeyPair::RSAMergeKeyPair(RSAKeyPair *kpub, RSAKeyPair *kprv)
{
   RSACopyPublicKey(pub, kpub->pub);
   RSACopyPrivateKey(prv, kprv->prv);
}

int RSAKeyPair::RSAKeyCompareThisPrivateToPublic(RSAKeyPair *icmp)
{
   EVP_PKEY *evp1 = EVP_PKEY_new();
   EVP_PKEY *evp2 = EVP_PKEY_new();
   EVP_PKEY_set1_RSA(evp1, prv);
   EVP_PKEY_set1_RSA(evp2, icmp->pub);
   int retval = EVP_PKEY_cmp(evp1, evp2);
   EVP_PKEY_free(evp1);
   EVP_PKEY_free(evp2);
   fSSLCheckError();
   switch (retval)
   {
      case -2:
         errprint("key compare operation not supported (EVP_PKEY_cmp())");
         break;
      case -1:
         errprint("tried to compare different key types (EVP_PKEY_cmp())");
         return 0;
      case 0:
         return 0;
      case 1:
         return 1;
   }
   return 0;
}
